# -*- coding: utf-8 -*-
import sys
import os
import progressbar
import tensorflow as tf
import numpy as np
sys.path.append(os.path.split(os.path.realpath(__file__))[0])

from mtcnn.mtcnn import mtcnn
from batch_processing import batch_processing
import cv2
import time

class mtcnn_batchprocess():
    def __init__(self):
        self.m_mtcnn = mtcnn(minsize = 100, maxsize = 0, threshold = [ 0.6, 0.7, 0.6], factor = 0.9)
        self.batch_processing = batch_processing()
        print 'mtcnn_batchprocess Init()!'

    def get_file_list(self, file_path, file_type = '.jpg'):
        path_list, file_list = self.batch_processing.get_file_list(file_path, file_type)
        return path_list, file_list
    
    def mtcnn_process(self, path_list, file_list, margin_value = 0, margin_ratio = 0.0,
                      resize_type = 0, resize_ratio = 1.0, controled_size = 0,
                      base_folder = '', save_folder = '', print_process = True, 
                      error_pass = False, error_moveFile = False):

#         margin_ratio > 0.0:  margin at ratio (margin_ratio)
#         margin_ratio = 0.0:  margin = margin_value              

#        if base_path != '' and save_folder_name != '' : save file to other path(save_path)
#             such as: image_path  = '/ai/oeasy_work/mtcnn_1124_1206/mtcnn_tf_caffe/mtcnn_tf/t1112.jpg'
#                      base_folder = 'mtcnn_1124_1206'
#                      save_folder = 'save'
#                 then:
#                      save_path = '/ai/oeasy_work/save/mtcnn_tf_caffe/mtcnn_tf/t1112.jpg'
#          else:
#             Replace the original file
        if margin_ratio >1.0:
            margin_ratio =1.0

        file_list_num = len(file_list)

        if print_process:
            p_progressbar = progressbar.ProgressBar(maxval = file_list_num)
            p_progressbar.start()

        for cur_index in xrange(file_list_num):
            time.sleep(0.02)
            if print_process:
                p_progressbar.update(cur_index+1)

            c_path = path_list[cur_index]
            c_file = file_list[cur_index]
            src_file = os.path.join(c_path, c_file)
            try:
                image_data = cv2.imread(src_file)
                img_size = image_data.shape
                bounding_boxes, points = self.m_mtcnn.detect_face(image_data, intput_type = 1)
                
                # There is only one person's face
                if type(bounding_boxes) is np.ndarray and bounding_boxes.shape[0] == 1:
                    c_rect = bounding_boxes[0]
                    if margin_ratio >0:
                        rect_w = int(c_rect[3]) - int(c_rect[1])
                        rect_h = int(c_rect[2]) - int(c_rect[0])
                        margin_w = rect_w * margin_ratio
                        margin_h = rect_h * margin_ratio
                    else:
                        margin_w = margin_value
                        margin_h = margin_value
                    
                    rect_margined = np.zeros(4, dtype=np.int32)
                    rect_margined[0] = np.maximum(c_rect[0] - margin_w, 0)
                    rect_margined[1] = np.maximum(c_rect[1] - margin_h, 0)
                    rect_margined[2] = np.minimum(c_rect[2] + margin_w, img_size[1])
                    rect_margined[3] = np.minimum(c_rect[3] + margin_h, img_size[0])
                    margined_data = image_data[rect_margined[1]:rect_margined[3],rect_margined[0]:rect_margined[2]]
                    dst_img = self.batch_processing.resize_image_data(margined_data, resize_type, resize_ratio, controled_size)
                    if base_folder != '' and save_folder != '':
                        c_save_folder = c_path.replace(base_folder, save_folder, 1)
                        if not os.path.exists(c_save_folder):
                            os.makedirs(c_save_folder)
                    else:
                        c_save_folder = c_path
                    dst_file = os.path.join(c_save_folder, c_file)
                    if not cv2.imwrite(dst_file, dst_img):
                        print ('no_permission : ' + src_file)
                else:
                    os.remove(src_file)
            except:
                print ('err_image : ' + src_file)
                if not error_pass:
                    return
                elif error_moveFile and os.path.exists(src_file):
                    os.remove(src_file)


        if print_process:
            p_progressbar.finish() 

